# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.
# Adapted from https://github.com/MARIO-Math-Reasoning/Super_MARIO
"""
rStar-Math MCTS 节点模块

这个模块实现了蒙特卡洛树搜索（MCTS）算法中的节点类。MCTSNode 继承自 BaseNode，
添加了 MCTS 算法特有的统计信息和计算方法。

主要功能：
1. 维护节点的访问次数和价值累积
2. 计算 Q 值（平均价值）
3. 实现 PUCT（Polynomial Upper Confidence Trees）公式
4. 支持递归价值回传（backpropagation）

"""
from __future__ import annotations

import numpy as np
from typing import Optional, Dict, Any, List, Type
from pydantic import BaseModel, PrivateAttr, field_validator
from .base_node import BaseNode


class MCTSNode(BaseNode):
    """
    MCTS 算法专用节点类
    
    继承自 BaseNode，添加了 MCTS 算法所需的统计信息和计算方法。
    实现了标准的 MCTS 节点功能，包括价值统计、PUCT 计算和递归更新。
    
    属性说明：
        c_puct: PUCT 公式中的探索常数，控制探索与利用的平衡
        inited: 节点是否已初始化标志
        __visit_count: 节点被访问的次数（私有属性）
        __value_sum: 节点价值的累积和（私有属性）
    
    MCTS 算法流程：
    1. Selection: 使用 PUCT 值选择最优子节点
    2. Expansion: 扩展选中的叶子节点
    3. Simulation: 模拟推理过程获得价值
    4. Backpropagation: 将价值递归传播到根节点
    """

    # MCTS 算法参数
    c_puct: float = 2        # PUCT 探索常数
    inited: bool = False     # 初始化标志

    # 私有统计属性（使用 Pydantic PrivateAttr）
    __visit_count: int = PrivateAttr(default=0)    # 访问次数
    __value_sum: float = PrivateAttr(default=0)    # 价值累积和

    def q_value(self) -> float:
        """
        计算节点的 Q 值（平均价值）
        
        Q 值是 MCTS 算法中评估节点质量的核心指标，表示该节点的平均回报。
        计算公式：Q = value_sum / visit_count
        
        返回:
            float: 节点的 Q 值，如果未被访问则返回 0
        """
        if self.__visit_count == 0:
            return 0
        return self.__value_sum / self.__visit_count

    def visit_count(self) -> int:
        """
        获取节点的访问次数
        
        返回:
            int: 节点被访问的次数
        """
        return self.__visit_count

    def update_visit_count(self, count: int) -> None:
        """
        更新节点的访问次数
        
        参数:
            count: 新的访问次数
        """
        self.__visit_count = count

    def update(self, value: float) -> None:
        """
        更新节点的统计信息
        
        增加访问次数并累积价值。如果是首次更新，则初始化节点的基础价值。
        这是 MCTS 算法中 backpropagation 阶段的核心操作。
        
        参数:
            value: 要累积的价值
        """
        if self.inited is False:
            self.inited = True
            self.value = value  # 设置节点的基础价值
        self.__visit_count += 1
        self.__value_sum += value

    def update_recursive(self, value: float, start_node: Type[BaseNode]) -> None:
        """
        递归更新节点价值（回传过程）
        
        从当前节点开始，向根节点方向递归传播价值更新。这是 MCTS 算法中
        backpropagation 阶段的实现，确保搜索过程中获得的信息能够影响
        整个搜索路径上的节点。
        
        参数:
            value: 要传播的价值（可以是列表或浮点数）
            start_node: 开始回传的起始节点，用于终止递归
        """
        # 处理价值格式（支持列表输入）
        if isinstance(value, list):
            value = float(value[0])
        
        # 更新当前节点
        self.update(value)
        
        # 递归终止条件：到达起始节点
        if self.tag == start_node.tag:
            return
        
        # 向父节点递归传播
        self.parent.update_recursive(value, start_node)

    def puct(self) -> float:
        """
        计算 PUCT（Polynomial Upper Confidence Trees）值
        
        PUCT 是 MCTS 算法中用于节点选择的核心公式，平衡了利用（exploitation）
        和探索（exploration）。公式为：
        PUCT = Q + c_puct * sqrt(ln(parent_visits) / child_visits)
        
        其中：
        - Q: 节点的平均价值（利用项）
        - c_puct: 探索常数
        - sqrt(...): 置信上界项（探索项）
        
        返回:
            float: 节点的 PUCT 值，用于选择最优子节点
        """
        # 根节点没有父节点，返回 0
        if not self.parent: 
            return 0
        
        # 计算 Q 值（利用项）
        q_value = self.q_value() if self.visit_count() > 0 else 0
        
        # 计算置信上界项（探索项）
        if self.parent.visit_count() == 0 or self.visit_count() == 0:
            u_value = 0  # 避免除零错误
        else:
            # UCB1 公式的置信上界项
            u_value = self.c_puct * np.sqrt(np.log(self.parent.visit_count()) / (self.visit_count()))
        
        # 返回 PUCT 值：利用项 + 探索项
        return q_value + u_value
        
